#include "mem/mmap.h"
#include "mem/pmem.h"
#include "mem/vmem.h"
#include "proc/cpu.h"
#include "lib/print.h"
#include "lib/str.h"
#include "memlayout.h"


//连续虚拟空间的复制(在uvm_copy_pgtbl中使用)
static void copy_range(pgtbl_t old, pgtbl_t new, uint64 begin, uint64 end)
{
    uint64 va, pa;
    int flags;
    pte_t* pte;

    for(va = begin; va < end; va += PGSIZE)
    {
        pte = vm_getpte(old, va, false);
        //懒分配
        //assert(pte != NULL, "uvm_copy_pgtbl: pte == NULL");
        //assert((*pte) & PTE_V, "uvm_copy_pgtbl: pte not valid");
        if (pte == NULL || !(*pte & PTE_V)) {
            continue;
        }

        pa = (uint64)PTE_TO_PA(*pte);
        page_node_t* page = (page_node_t*)pa; // 将物理地址转换为 page_node_t*
        page->refNum++;
        *pte = (*pte & ~PTE_W) | PTE_COW;
        flags = (int)PTE_FLAGS(*pte);

        // page = (uint64)pmem_alloc(false);
        // memmove((char*)page, (const char*)pa, PGSIZE);
        vm_mappages(new, va, pa, PGSIZE, flags);
    }
}

// 两个 mmap_region 区域合并
// 保留一个 释放一个 不操作 next 指针
// 在uvm_munmap里使用
static void mmap_merge(mmap_region_t* mmap_1, mmap_region_t* mmap_2, bool keep_mmap_1)
{
    // 确保有效和紧临
    assert(mmap_1 != NULL && mmap_2 != NULL, "mmap_merge: NULL");
    assert(mmap_1->begin + mmap_1->npages * PGSIZE == mmap_2->begin, "mmap_merge: check fail");
    
    // merge
    if(keep_mmap_1) {
        mmap_1->npages += mmap_2->npages;
        mmap_region_free(mmap_2);
    } else {
        mmap_2->begin -= mmap_1->npages * PGSIZE;
        mmap_2->npages += mmap_1->npages;
        mmap_region_free(mmap_1);
    }
}

// 打印以 mmap 为首的 mmap 链
// for debug
void uvm_show_mmaplist(mmap_region_t* mmap)
{
    mmap_region_t* tmp = mmap;
    printf("\nmmap allocable area:\n");
    if(tmp == NULL)
        printf("NULL\n");
    while(tmp != NULL) {
        printf("allocable region: %p ~ %p\n", tmp->begin, tmp->begin + tmp->npages * PGSIZE);
        tmp = tmp->next;
    }
}

// 递归释放 页表占用的物理页 和 页表管理的物理页
// ps: 顶级页表level = 3, level = 0 说明是页表管理的物理页
static void destroy_pgtbl(pgtbl_t pgtbl, uint32 level)
{
    pte_t page_entry;   // 当前页表项
    pgtbl_t child_table; // 子页表指针

    // 检查页表级别
    assert(level <= 3, "free_page_table: Invalid level");

    if (level == 0) 
        goto clean;  // 物理页，无需进一步处理

    // 遍历页表中的每一项
    for (int i = 0; i < PGSIZE / sizeof(page_entry); i++) {
        page_entry = pgtbl[i];
        if (page_entry & PTE_V) {  // 当前页表项有效
            if (level > 1 && !(PTE_CHECK(page_entry))) {
                // 页表项错误，未映射为页表项
                printf("page_entry = %p\n", page_entry);
                panic("destroy_pgtbl: Invalid page entry");
            }
            child_table = (pgtbl_t)PTE_TO_PA(page_entry);  // 获取子页表的物理地址
            destroy_pgtbl(child_table, level - 1);  // 递归销毁子页表
            pgtbl[i] = 0;  // 清空当前页表项
        }
    }

    clean:
        // 根据页表级别释放内存
        if (level > 0)
            pmem_free((uint64)pgtbl, true);
        else
            pmem_free((uint64)pgtbl, false);
}

// 页表销毁：trapframe 和 trampoline 单独处理
void uvm_destroy_pgtbl(pgtbl_t pgtbl)
{
    vm_unmappages(pgtbl, TRAPFRAME, PGSIZE, true);
    vm_unmappages(pgtbl, TRAMPOLINE, PGSIZE, false);
    destroy_pgtbl(pgtbl, 3);
}


// 拷贝页表 (拷贝并不包括trapframe 和 trampoline)
void uvm_copy_pgtbl(pgtbl_t old, pgtbl_t new, uint64 heap_top, uint32 ustack_pages, mmap_region_t* mmap)
{
    /* step-1: USER_BASE ~ heap_top */
    copy_range(old, new, USER_BASE, heap_top);

    /* step-2: ustack */
    copy_range(old, new, TRAPFRAME - ustack_pages * PGSIZE, TRAPFRAME);

    /* step-3: mmap_region */
    uint64 alloc_start = MMAP_BEGIN;
    for (mmap_region_t *region = mmap; region != NULL; region = region->next)
    {
        while (alloc_start<region->begin)
        {
            copy_range(old, new, alloc_start, alloc_start + PGSIZE);
            alloc_start += PGSIZE;
        }
    }
}

// 在用户页表和进程mmap链里 新增mmap区域 [begin, begin + npages * PGSIZE)
// 页面权限为perm
void uvm_mmap(uint64 begin, uint32 npages, int perm)
{
    // 如果需要映射的页面数量为 0，直接返回
    if (npages == 0) return;  

    // 确保起始地址对齐
    assert(begin % PGSIZE == 0, "uvm_mmap: start not aligned");  

    struct proc *current_process = myproc(); 
    uint64 end = begin + npages * PGSIZE;
    mmap_region_t *current_region = current_process->mmap; 

    // 遍历 mmap 链表，查找合适的插入位置
    while (current_region) {
        // 情况 1: 请求的内存区域位于当前区域的中间，拆分当前区域
        if (begin > current_region->begin && end < current_region->begin + current_region->npages * PGSIZE) {
            mmap_region_t *new_region = pmem_alloc(false);
            new_region->begin = begin + npages * PGSIZE;  
            current_region->npages = (begin - current_region->begin) / PGSIZE;  
            new_region->npages = current_region->npages - current_region->npages - npages;  

            new_region->next = current_region->next;  
            current_region->next = new_region;  
            break;
        }// 情况 2: 请求的内存区域从当前区域的起始位置开始
        else if (begin == current_region->begin) {
            current_region->begin = begin + npages * PGSIZE; 
            current_region->npages -= npages; 
            break;
        }
        else if (end == current_region->begin + current_region->npages * PGSIZE) {
            // 情况 3: 请求的内存区域恰好接在当前区域的末尾
            current_region->npages -= npages;  
            break;
        }

        current_region = current_region->next; 
    }

    // 修改页表，分配物理页并映射到虚拟地址
    for (uint32 i = 0; i < npages; i++) {
        char *physical_page = pmem_alloc(false);  // 分配新的物理页
        vm_mappages(current_process->pgtbl, begin + i * PGSIZE, (uint64)physical_page, PGSIZE, perm);  // 映射物理页到虚拟地址
    }
}

// 在用户页表和进程mmap链里释放mmap区域 [begin, begin + npages * PGSIZE)
void uvm_munmap(uint64 begin, uint32 npages)
{
    if (npages == 0) return;
    assert(begin % PGSIZE == 0, "uvm_munmap: begin not aligned");
    struct proc *proc = myproc();

    uint64 end = begin + npages * PGSIZE;

    // new mmap_region 的产生
    mmap_region_t *new_region = pmem_alloc(false);
    new_region->begin = begin;
    new_region->npages = npages;
    new_region->next = NULL;

    mmap_region_t *mp = proc->mmap;

    // 尝试合并 mmap_region
    while (mp)
    {
        // 在链表第一个节点前面且不能合并
        if (end<mp->begin)
        {
            proc->mmap = new_region;
            new_region->next = mp;
            break;
        }
        //在整个链表最后面
        else if(mp->next==NULL){
            mp->next = new_region;
            new_region->next = NULL;
        }
        // 接在上一个mmap_region后面
        else if (mp->begin+mp->npages*PGSIZE==begin)
        {
            mmap_merge(mp,new_region,1);
            if (end == mp->next->begin)
            {
                mmap_merge(mp,mp->next,1);
                mp->next = mp->next->next;
            }
            break;
        }
        // 接在下一个mmap_region前面，此时与上一个mmap_region不相连
        else if (end==mp->begin)
        {
            mmap_merge(new_region,mp,0);
            break;
        }
        // 在两个mmap_region之间再插入一个mmap_region且不能合并
        else if (begin>mp->begin+mp->npages+PGSIZE && end<mp->next->begin)
        {
            new_region->next = mp->next;
            mp->next = new_region;
            break;
        }

        mp = mp->next;
    }

    // 页表释放
    for (uint64 addr = begin; addr < end; addr += PGSIZE)
    {
        pte_t *pte = vm_getpte(proc->pgtbl, addr, 0);
        if (pte && (*pte & PTE_V))
        {
            vm_unmappages(proc->pgtbl, addr, PGSIZE, 1);
        }
    }
}

// 用户堆空间增加, 返回新的堆顶地址 (注意栈顶最大值限制)
// 在这里无需修正 p->heap_top
uint64 uvm_heap_grow_lazy(pgtbl_t pgtbl, uint64 heap_top, uint32 len, int flags)
{
    uint64 new_heap_top = heap_top + len;
    // 假设堆增长方向是从低地址到高地址，进行逐页分配
    uint64 va = heap_top;
    uint64 end = new_heap_top;
    //char *mem;
    // 只在页表中创建映射，不立即分配内存
    while (va < end){
        vm_mappages(pgtbl, va, 0, PGSIZE, PTE_U | flags);
        va += PGSIZE;
    }

    return new_heap_top;
}

// 用户堆空间增加, 返回新的堆顶地址 (注意栈顶最大值限制)
// 在这里无需修正 p->heap_top
uint64 uvm_heap_grow(pgtbl_t pgtbl, uint64 heap_top, uint32 len, int flags)
{
    uint64 new_heap_top = heap_top + len;
    // 假设堆增长方向是从低地址到高地址，进行逐页分配
    uint64 va = heap_top;
    uint64 end = new_heap_top;
    char *mem;
    // 只在页表中创建映射，不立即分配内存
    while (va < end){
        mem=pmem_alloc(false);
        if(mem==NULL){
            printf("Error: pmem_alloc failed at virtual address 0x%lx\n", va);
            panic("uvm_heap_grow: pmem_alloc failed");
        }
        vm_mappages(pgtbl, va, (uint64)mem, PGSIZE, PTE_U | flags);
        va += PGSIZE;
    }

    return new_heap_top;
}



// 用户堆空间减少, 返回新的堆顶地址
// 在这里无需修正 p->heap_top
uint64 uvm_heap_ungrow(pgtbl_t pgtbl, uint64 heap_top, uint32 len)
{
    uint64 new_heap_top = heap_top - len;

    vm_unmappages(pgtbl, new_heap_top, len, true);

    return new_heap_top;
}

// 用户态地址空间[src, src+len) 拷贝至 内核态地址空间[dst, dst+len)
// 注意: src dst 不一定是 page-aligned
void uvm_copyin(pgtbl_t pgtbl, uint64 dst, uint64 src, uint32 len)
{
    uint64 n, va0, pa0;
    
    while(len > 0) {
        // 将 src 对齐到页边界
        va0 = ALIGN_DOWN(src,PGSIZE);
        
        // 获取用户空间地址 va0 对应的页表项
        pte_t *p = vm_getpte(pgtbl, va0, false); // false 表示不申请页
      
        if (p == NULL) {
            // 如果页表项为空，说明用户虚拟地址无效
            // 可以打印错误信息来帮助调试
        panic("Invalid user virtual address");
        }

        pa0 = PTE_TO_PA(*p);  // 获取物理地址

        if(pa0 == 0) {
            // 如果物理地址为 0，说明该页面未映射或访问无效
            // 此时也可以打印调试信息来帮助定位错误
            printf("Error: Invalid mapping at virtual address %p\n", va0);
            return;  // 返回或处理错误
        }

        // 计算当前页可以拷贝的数据量
        n = PGSIZE - (src - va0);  // 当前页剩余的可用字节数
        
        // 如果剩余字节数大于需要拷贝的字节数，截断拷贝长度
        if(n > len) {
            n = len;
        }

        // 执行数据拷贝，拷贝 src -> dst
        memmove((void*)dst, (void*)(pa0 + (src - va0)), n);
        
        // 更新剩余长度
        len -= n;
        // 更新目标地址和源地址
        dst += n;
        src += n;
        
        // 如果当前页已经拷贝完，则处理下一页
        if(src == va0 + PGSIZE) {
            src = va0 + PGSIZE;
        }
    }

}

// 内核态地址空间[src, src+len） 拷贝至 用户态地址空间[dst, dst+len)
void uvm_copyout(pgtbl_t pgtbl, uint64 dst, uint64 src, uint32 len)
{
    uint64 n, va0, pa0;
    while(len > 0) {
        // 将 dst 对齐到页边界
        va0 = ALIGN_DOWN(dst, PGSIZE);
        // 获取内核空间地址 va0 对应的页表项
        pte_t *p = vm_getpte(pgtbl, va0, false);  // false 表示不申请页
        if (p == NULL) {
            // 如果页表项为空，说明地址无效
            panic("Invalid kernel virtual address");
        }
        pa0 = PTE_TO_PA(*p);  // 获取物理地址
        if (pa0 == 0) {
            // 如果物理地址为 0，说明该页面未映射或访问无效
            //printf("Error: Invalid mapping at virtual address %p\n", va0);
            return;  // 返回或处理错误
        }
        // 计算当前页可以拷贝的数据量
        n = PGSIZE - (dst - va0);  // 当前页剩余的可用字节数
        // 如果剩余字节数大于需要拷贝的字节数，截断拷贝长度
        if (n > len) {
            n = len;
        }
        // 执行数据拷贝，拷贝 src -> dst
        memmove((void *)(pa0 + (dst - va0)), (void *)src, n);
        // 更新剩余长度
        len -= n;
        // 更新目标地址和源地址
        dst += n;
        src += n;
        // 如果当前页已经拷贝完，则处理下一页
        if (dst == va0 + PGSIZE) {
            dst = va0 + PGSIZE;
        }
    }

}

// 用户态字符串拷贝到内核态
// 最多拷贝maxlen字节, 中途遇到'\0'则终止
// 注意: src dst 不一定是 page-aligned
void uvm_copyin_str(pgtbl_t pgtbl, uint64 dst, uint64 src, uint32 maxlen)
{
    uint64 n, va0, pa0;
    int got_null = 0;

    while (maxlen > 0 && !got_null) {
        // 将 src 对齐到页边界
        va0 = ALIGN_DOWN(src, PGSIZE); // 使用 ALIGN_DOWN 宏对齐地址
        
        // 查找用户空间页表项
        pte_t *pte = vm_getpte(pgtbl, va0, 0);  // 获取对应的页表项
        pa0 = PTE_TO_PA(*pte);  // 获取物理地址

        if (pa0 == 0) {
            return; // 如果物理地址为 0，表示无效映射，返回
        }

        // 计算当前页可以拷贝的字节数
        n = PGSIZE - (src - va0);  // 当前页剩余可拷贝的字节数
        if (n > maxlen) {
            n = maxlen;  // 如果剩余的字节数大于 maxlen，限制为 maxlen
        }

        // 拷贝数据，从 src 到 dst
        char *src_ptr = (char *)(pa0 + (src - va0));  // 将 src 转换为物理地址对应的虚拟地址
        char *dst_ptr = (char *)dst;  // 将 dst 转换为目标地址

        while (n > 0) {
            *dst_ptr = *src_ptr;  // 拷贝一个字节
            if (*src_ptr == '\0') {  // 如果遇到 '\0'，终止复制
                got_null = 1;
                break;
            }
            n--;       // 更新剩余字节数
            maxlen--;  // 更新最大字节数
            src_ptr++; // 移动到下一个字节
            dst_ptr++; // 移动目标地址
        }

        // 更新 src 和 dst 地址，准备拷贝下一页的数据
        src += (n);  // 将 src 地址移动 n 个字节
        dst += (n);  // 将 dst 地址移动 n 个字节
    }

}